package com.hnit.service;

import com.hnit.config.RabbitMQConfig;
import org.springframework.amqp.rabbit.annotation.RabbitListener;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;

@Service
public class VectorizationConsumer {

    private final VectorStoreService vectorStoreService;

    @Autowired
    public VectorizationConsumer(VectorStoreService vectorStoreService) {
        this.vectorStoreService = vectorStoreService;
    }

    @RabbitListener(queues = RabbitMQConfig.VECTORIZATION_QUEUE)
    public void processVectorizationTask(Map<String, Object> taskData) {
        System.out.println("Consumer started processing message from queue: " + RabbitMQConfig.VECTORIZATION_QUEUE);
        try {
            String taskId = (String) taskData.get("taskId");
            System.out.println("Received task: " + taskId + ", data: " + taskData);

            // 处理单个JSON数据
            if (taskData.containsKey("title") || taskData.containsKey("content")) {
                System.out.println("Processing single JSON data for task: " + taskId);
                int segmentCount = vectorStoreService.addJsonToVectorStore(taskData);
                System.out.println("Task " + taskId + " processed: " + segmentCount + " segments");
            }
            // 处理批量JSON数据
            else if (taskData.containsKey("dataList")) {
                List<Map<String, Object>> jsonDataList = (List<Map<String, Object>>) taskData.get("dataList");
                System.out.println("Processing batch JSON data for task: " + taskId);
                int totalSegmentCount = vectorStoreService.addJsonListToVectorStore(jsonDataList);
                System.out.println("Batch task " + taskId + " processed: " + totalSegmentCount + " segments");
            }
            // 处理文本数据
            else if (taskData.containsKey("content")) {
                String content = (String) taskData.get("content");
                System.out.println("Processing text data for task: " + taskId);
                int segmentCount = vectorStoreService.addTextToVectorStore(content);
                System.out.println("Text task " + taskId + " processed: " + segmentCount + " segments");
            } else {
                System.err.println("Invalid task data for taskId " + taskId + ": " + taskData);
            }
        } catch (Exception e) {
            System.err.println("Vectorization task failed for taskId " + taskData.get("taskId") + ": " + e.getMessage());
            e.printStackTrace();
        }
    }
}